'''
yolov8 实例分割推理
'''
import glob
import time

import cv2
import numpy as np
import torch

from ultralytics import YOLO # v8.1
from line_profiler import LineProfiler # 逐行计算运行时间
ctime = LineProfiler()

def get_yolov8_seg_predict(model_path = 'yolov8n-seg.pt'):
    # model_path = 'yolov8n.pt'
    if isinstance(model_path, YOLO):
        model = model_path
    else:
        model = YOLO(model_path)
    # device = 'cpu'
    device = '0' if torch.cuda.is_available() else 'cpu'
    model.predict(np.zeros((640,640,3), np.uint8), imgsz=640, device=device, retina_masks=True) # 预热
    # model.predictor(imgsz=320, conf=0.5)
    # @ctime
    def yolov8_seg_predict(img, imgsz=640, realtime_conf_thres=0.35, iou_thres=0.45, roi=None, pad=0,
                           isShowMask=False, show_img=None, thickness_rate=0.001, isMask2Xy=False, isDraw=True):
        '''
        推理单张图片， 可roi检测
        show_img, out = seg_v8_step1xjg(img0, realtime_conf_thres=0.5, iou_thres=0.45,
                                        roi=fix_bd_step1['xiangjigai'], pad=20, isShowMask=False)
        Args:
            img: 2048*3072*3 原图
            realtime_conf_thres:
            iou_thres:
            roi: [xyxy] 是否roi
            pad: 填充 px
            isShowMask: 是否画mask
            show_img : 可外部传入show_img, 为None则copy原图
            isMask2Xy: mask是否转为轮廓(numpy 1*n*2)，加快速度
            isDraw: 是否画图
        Returns:
            show_img:
            out: [[xyxy],cls_ind, clsname,conf,mask] [list(1*4),int,str,float,numpy(1*h*w)]

        '''
        H, W = img.shape[:2]
        if roi is not None:
            x1, y1, x2, y2 = roi
            roi = [int(max(x1 - pad, 0)), int(max(y1 - pad, 0)), int(min(x2 + pad, W)), int(min(y2 + pad, H))]

        img2 = img[roi[1]:roi[3],roi[0]:roi[2],...] if roi is not None else img
        t0 = time.time()
        results = model.predict(img2, imgsz=imgsz, conf=realtime_conf_thres, iou=iou_thres, device=device, retina_masks=True)  # retina_masks 高分辨率mask
        print(f'3333333333333333333 {time.time()-t0}')
        result = results[0]  # one img
        boxes = result.boxes  # Boxes object for bounding box outputs
        masks = result.masks  # Masks object for segmentation masks outputs  n*448*640
        # keypoints = result.keypoints  # Keypoints object for pose outputs
        # probs = result.probs  # Probs object for classification outputs
        # result.show()  # display to screen
        # result.save(filename='result.jpg')  # save to disk
        boxes = boxes.cpu().numpy()
        # if boxes.shape[0] == 0:
            # return img.copy(), []
        if masks is not None:
            # masks = masks.cpu().numpy()  # 当无目标时masks = None
            # print(masks.shape)
            if not isMask2Xy:
                masks = masks.cpu().numpy() # numpy 1*h*w  0.0 or 1.0
            else:
                masks = masks.xy # [n*2,n*2] 18ms
        out = []
        for i in range(boxes.shape[0]):
            xyxy = boxes.xyxy[i, :].tolist()
            conf = boxes.conf[i]
            cls_ind = int(boxes.cls[i])
            cls_name = result.names[cls_ind]
            # if not isMask2Xy:
            #     mask = masks[i].data.cpu().numpy()  # numpy 1*h*w  0.0 or 1.0
            # else:
            #     # mask = masks[i].xy[0].astype(np.int32) #  numpy n*2 14ms
            #     mask = masks[i].xy[0]  # numpy n*2 14ms
            mask = masks[i].data if not isMask2Xy else masks[i].reshape(1, -1, 2) # 1*h*w
            if roi is not None:  # 缩放回原图
                xyxy[0] += roi[0]
                xyxy[1] += roi[1]
                xyxy[2] += roi[0]
                xyxy[3] += roi[1]
                if not isMask2Xy:
                    mask = np.pad(mask, ((0, 0), (roi[1], H - roi[3]), (roi[0], W - roi[2])), 'constant',
                                  constant_values=(0, 0))  # 1*h*w 填充 耗时40ms
                else:
                    mask = mask + roi[:2]
            out.append([xyxy, cls_ind, cls_name, conf, mask])
        print(f'444444444444444444 {time.time() - t0}')

        # draw
        if not isDraw: # 是否画图
            show_img = None
        else:
            if show_img is None:
                show_img = img.copy()
            colors = [(255, 0, 0), (0, 255, 0), (0, 0, 255), (0, 255, 255), (255, 255, 0), (255, 0, 255)]
            img_diag = np.sqrt(show_img.shape[0] ** 2 + show_img.shape[1] ** 2)
            print(img_diag)
            for ind, (xyxy, cls_ind, cls_name, conf, mask) in enumerate(out):
                if isShowMask:
                    if not isMask2Xy:
                        print(mask.shape)
                        # mask = mask.repeat(3, axis=0).transpose(1, 2, 0)  # 3*h*w 彩色  耗时
                        mask = mask.reshape(mask.shape[1], mask.shape[2],1).repeat(3, axis=2)  # hw1 to 3*h*w 彩色  耗时
                        # mask = show_img
                        # show_img = (show_img + mask * colors[cls_ind%len(colors)] * 0.4).clip(0, 255).astype(np.uint8) # 同类别同色
                        show_img = (show_img + mask * colors[ind % len(colors)] * 0.4).clip(0, 255).astype(np.uint8) # 值传递
                        # show_img = show_img.getUMat()
                    else:
                        cv2.polylines(show_img, [mask.reshape(-1, 1, 2).astype(np.int32)], False, (0, 255, 0), max(1, int(img_diag * thickness_rate)))

                px1, py1, px2, py2 = int(xyxy[0]), int(xyxy[1]), int(xyxy[2]), int(xyxy[3])
                print(show_img.shape)
                cv2.putText(show_img, f'{str(np.round(conf,2))}_{cls_ind}_{cls_name}', (px1, py1),
                            cv2.FONT_HERSHEY_COMPLEX_SMALL, max(1, int(img_diag * thickness_rate*0.5)), (0, 255, 0),
                            max(1, int(img_diag * 0.0005)))
                # cv2.rectangle(show_img)
                cv2.rectangle(show_img, (px1, py1), (px2, py2), (0, 255, 0), max(1, int(img_diag * thickness_rate)))
            if roi is not None:
                cv2.rectangle(show_img, (roi[0], roi[1]), (roi[2], roi[3]), (255, 0, 0), max(1, int(img_diag * thickness_rate)))
        print(f'555555555555555 {time.time() - t0}')
        return show_img, out

    return yolov8_seg_predict



def ttest_get_detect_onepic():
    img = cv2.imread(r"D:\data\231207huoni\trainV8Seg_cable\add_imgs\20240104\Image_20240104150413321.jpg")
    model_path = r"D:\data\231207huoni\trainV8Seg_cable\models\640_cable\weights\best.pt"
    detect_onepic = get_yolov8_seg_predict(model_path)
    img_show, out = detect_onepic(img,0.25,roi=[100,100,2500,1500], isShowMask=True, isMask2Xy=True)
    ctime.print_stats()
    cv2.imwrite(r'D:\data\231207huoni\test_data\1.jpg', img_show)
    # print(out)

def ttest_dir(model_path = '', img_glob = '', save_dir = ''):
    # from ultralytics import YOLO
    # test_dir = r"D:\data\231215安全带\trainV8Det_flball_blue\format_data\images\train"
    # # Load a pretrained YOLOv8n model
    # model_path = r"D:\data\231215安全带\trainV8Det_flball_blue\models\train7\weights\best.pt"
    # model = YOLO(model_path)
    # # Run inference on 'bus.jpg' with arguments
    # model.predict(test_dir, save=True, imgsz=640, conf=0.5, show=True)

    ls = glob.glob(img_glob)
    detect_onepic = get_yolov8_seg_predict(model_path)
    for ind, i in enumerate(ls):
        img = cv2.imread(i)
        # print(i)
        # print(img)
        # try:
        img_show, out = detect_onepic(img, 0.25, isShowMask= False)
        cv2.imwrite(f'{save_dir}/{ind}.jpg', img_show)
        # except:
        #     print(f'error {i}')

if __name__ == '__main__':
    ttest_get_detect_onepic()
    # root_path = r'D:\data\231207huoni'
    # # root_path = r'/home/ps/zhangxiancai/data/231207huoni/'
    # ttest_dir(model_path = rf"{root_path}/trainV8Seg_screw/models/yolov8mSeg_640_screw/weights/best.pt",
    #           img_glob = rf'{root_path}/trainV8Seg_screw/add_imgs/*/*.jpg',
    #           save_dir = rf"{root_path}/trainV8Seg_screw/predict")
